import argparse
import os
import shutil
import numpy as np
from tqdm import tqdm
import mmcv


def parse_config():
    parser = argparse.ArgumentParser(description='args')
    parser.add_argument('--data_paths', type=list, default=[
                                                            '/cv/yc/DSGN2/data/0915_dataset_ww',
                                                            '/cv/yc/DSGN2/data/ww_old_dataset'
                                                            ])
    parser.add_argument('--out_dir', type = str, default='/cv/yc/DSGN2/data/ww')
    parser.add_argument('--split', type=str, default='training', choices=['training', 'testing'])
    parser.add_argument('--dataset', type=str, default='train', choices=['train', 'val', 'test'])
    parser.add_argument('--kind', type=str, default='concat', choices=['image', 'lidar', 'concat'])
    parser.add_argument('--all_cameras', type=list, default=[0,2,3,4,6,7])
    parser.add_argument('--valid-cameras', type=list, default=[[0,2,3,4,6,7]] * 13)
    parser.add_argument('--image_size', type=str, default=[540, 960, 3])
    args = parser.parse_args()
    return args



if __name__ == '__main__':
    args = parse_config()
    data_paths = args.data_paths
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)
    tar_imageset_path = os.path.join(args.out_dir, 'ImageSets')
    if not os.path.exists(tar_imageset_path):
        os.makedirs(tar_imageset_path)
    # tar_data_path = os.path.join(args.out_dir, args.split)
    for idx, data_path in enumerate(data_paths):
        origin_txt_file = os.path.join(data_path, 'ImageSets', args.dataset + '.txt')
        if args.kind != 'image':
            ori_velodyne_dir = os.path.join(data_path, args.split, 'velodyne')
            ori_label_dir = os.path.join(data_path, args.split, 'label_2')
            # ori_ground_split_dir = os.path.join(data_path, args.split, 'ground_split')
        if args.kind != 'lidar':
            ori_calib_dir = os.path.join(data_path, args.split, 'calib')
            ori_images_dir = os.path.join(data_path, args.split, 'image_2')

        with open(origin_txt_file, 'r') as f:
            lines = f.readlines()
        lines = [line.strip('\n') for line in lines]
        for i in tqdm(range(len(lines))):
            timestamp = lines[i]
            with open(os.path.join(tar_imageset_path, args.dataset + '.txt'), 'a') as f:
                f.write(timestamp + '\n')
            if args.kind != 'image':
                ori_velodyne_path = os.path.join(ori_velodyne_dir, timestamp + '.bin')
                tar_velodyne_path = ori_velodyne_path.replace(data_path, args.out_dir)
                if not os.path.exists(os.path.dirname(tar_velodyne_path)):
                    os.makedirs(os.path.dirname(tar_velodyne_path))
                shutil.copy(ori_velodyne_path, tar_velodyne_path)
                ori_label_path = os.path.join(ori_label_dir, timestamp + '.txt')
                tar_label_path = ori_label_path.replace(data_path, args.out_dir)
            #     tar_label_path = tar_label_path.replace('label_2', 'label_2')
                if not os.path.exists(os.path.dirname(tar_label_path)):
                    os.makedirs(os.path.dirname(tar_label_path))
                shutil.copy(ori_label_path, tar_label_path)
            if args.kind != 'lidar':
                calib_cameras = os.listdir(ori_calib_dir)
                image_cameras = os.listdir(ori_images_dir)
                if image_cameras != calib_cameras:
                    valid_cameras = [camera for camera in calib_cameras if camera in image_cameras]
                else:
                    valid_cameras = calib_cameras
                valid_cameras = [camera for camera in valid_cameras if int(camera[-1]) in args.valid_cameras[idx]]
                for camera in args.all_cameras:
                    ori_calib_path = os.path.join(ori_calib_dir, f'camera{camera}', timestamp + '.txt')
                    tar_calib_path = ori_calib_path.replace(data_path, args.out_dir)
                    if not os.path.exists(os.path.dirname(tar_calib_path)):
                        os.makedirs(os.path.dirname(tar_calib_path))
                    shutil.copy(ori_calib_path, tar_calib_path)
                    #
                    #
                    ori_image_path = os.path.join(ori_images_dir, f'camera{camera}', timestamp + '.jpg')
                    tar_image_path = ori_image_path.replace(data_path, args.out_dir)
                    if f'camera{camera}' in valid_cameras:
                        if not os.path.exists(os.path.dirname(tar_image_path)):
                            os.makedirs(os.path.dirname(tar_image_path))
                        shutil.copy(ori_image_path, tar_image_path)
                    else:
                        img = np.zeros(args.image_size)
                        mmcv.imwrite(img,tar_image_path)


